#include <stdio.h>
#include <stdlib.h>
#include <mpi.h>
#include <math.h>

void verify(double *y_cpu, double *y, int size) {
    double eps = 1e-9;
    int error_cnt = 0;
    for (int i = 0; i < size; i++) {
        if (fabs(y_cpu[i] - y[i]) > eps) {
            error_cnt++;
        }
    }
    if(error_cnt > 0) {
        printf("Error! Total %d errors found!\n", error_cnt);
    } else {
        printf("Verification passed!\n");
    }
}

void mat_mul(double *a, double *b, double *c, int m, int n, int k) {
    for(int i = 0; i < m; i++) {
        for(int j = 0; j < k; j++) {
            c[i * k + j] = 0.0;
            for(int l = 0; l < n; l++) {
                c[i * k + j] += a[i * n + l] * b[l * k + j];
            }
        }
    }
}

int main(int argc, char **argv)
{
    int m = 0, n = 0, k = 0, myid, numprocs, srow = 0, i, p, srow_last = 0;
    double *M, *N, *P_serial, *P_mpi;

    double time_start_serial, time_end_serial, time_serial;
    double time_start_mpi, time_end_mpi, time_mpi;

    FILE *fp; 
    FILE *fp_out; 
    
    MPI_Init(&argc, &argv);
    MPI_Comm_size(MPI_COMM_WORLD, &numprocs);
    MPI_Comm_rank(MPI_COMM_WORLD, &myid);
    printf("I am proc %d\n", myid);
    
    if (myid == 0)
    {
        char s[20] = "input1.txt";
        fp = fopen(s, "r");
        fscanf(fp, "%d,%d,%d", &m, &n, &k);
        fclose(fp);
        M = (double *)malloc(m * n * sizeof(double));
        N = (double *)malloc(n * k * sizeof(double));
        P_serial = (double *)malloc(m * k * sizeof(double));
        P_mpi = (double *)malloc(m * k * sizeof(double));
        for (i = 0; i < m * n; i++)
            M[i] = (double)rand() / RAND_MAX;
        for (i = 0; i < n * k; i++)
            N[i] = (double)rand() / RAND_MAX;

        time_start_serial = MPI_Wtime();
        mat_mul(M, N, P_serial, m, n, k);
        time_end_serial = MPI_Wtime();
        time_serial = time_end_serial - time_start_serial;
        printf("Serial time: %.2f seconds\n", time_serial);
        char s_out[20] = "output1.txt";
        fp_out = fopen(s_out, "w");
    }
    if (numprocs == 1) {
        time_start_mpi = MPI_Wtime();
        mat_mul(M, N, P_mpi, m, n, k);
        time_end_mpi = MPI_Wtime();
        time_mpi = time_end_mpi - time_start_mpi;
        printf("MPI time with 1 process: %.2f seconds\n", time_mpi);
    }
    else{
        if(myid == 0){
            time_start_mpi = MPI_Wtime();
        }
        MPI_Bcast(&m, 1, MPI_INT, 0, MPI_COMM_WORLD);
        MPI_Bcast(&n, 1, MPI_INT, 0, MPI_COMM_WORLD);
        MPI_Bcast(&k, 1, MPI_INT, 0, MPI_COMM_WORLD);
        if (myid != 0) {
            N = (double *)malloc(n * k * sizeof(double));
        }
        MPI_Bcast(N, n * k, MPI_DOUBLE, 0, MPI_COMM_WORLD);
        
        srow = m / numprocs;
        srow_last = m - srow * (numprocs - 1);
        if (myid == numprocs - 1)
            srow = srow_last;

        if (myid == 0)
        {
            for (i = 1; i < numprocs - 1; i++)
                MPI_Send(M + i * srow * n, srow * n, MPI_DOUBLE, i, 0, MPI_COMM_WORLD);
            MPI_Send(M + i * srow * n, srow_last * n, MPI_DOUBLE, numprocs - 1, 0, MPI_COMM_WORLD);

            mat_mul(M, N, P_mpi, srow, n, k);

            for (i = 1; i < numprocs - 1; i++)
                MPI_Recv(P_mpi + i * srow * k, srow * k, MPI_DOUBLE, i, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
            MPI_Recv(P_mpi + i * srow * k, srow_last * k, MPI_DOUBLE, numprocs - 1, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);

            time_end_mpi = MPI_Wtime();
            time_mpi = time_end_mpi - time_start_mpi;
            printf("MPI time with %d processes: %.2f seconds\n", numprocs, time_mpi);

            verify(P_serial, P_mpi, m * k);
            fprintf(fp_out, "%.2f,%.2f", time_serial, time_mpi);
        }
        else
        {
            M = (double *)malloc(srow * n * sizeof(double));
            P_mpi = (double *)malloc(srow * k * sizeof(double));

            MPI_Recv(M, srow * n, MPI_DOUBLE, 0, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
            mat_mul(M, N, P_mpi, srow, n, k);
            MPI_Send(P_mpi, srow * k, MPI_DOUBLE, 0, 0, MPI_COMM_WORLD);
        }
    }
    if (myid == 0) {
        free(M);
        free(N);
        free(P_serial);
        free(P_mpi);
    } else {
        free(M);
        free(N);
        free(P_mpi);
    }
    MPI_Finalize();
    return 0;
}
